/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "aicpu_kernel_builder.h"

#include "cce/taskdown_common.hpp"
#include "cce/fwk_adpt_struct.h"
#include "aicpu/common/aicpu_task_struct.h"

#include "util/constant.h"
#include "error_code/error_code.h"
#include "util/log.h"
#include "util/util.h"
#include "config/ops_json_file.h"

namespace {
const std::string kDefaultKernelSo = "libcpu_kernels.so";
const std::string kDefaultFunctionName = "RunCpuKernel";

static uint64_t g_aicpu_kernel_id = 0;
static uint64_t g_aicpu_session_id = 0;

uint64_t GenerateUniqueKernelId()
{
    if (g_aicpu_kernel_id == ULLONG_MAX) {
        g_aicpu_kernel_id = 0;
    }
    return g_aicpu_kernel_id++;
}

uint64_t GenerateUniqueSessionId()
{
    if (g_aicpu_session_id == ULLONG_MAX) {
        g_aicpu_session_id = 0;
    }
    return g_aicpu_session_id++;
}
}

namespace aicpu {
KernelBuilderPtr AicpuKernelBuilder::instance_ = nullptr;

inline KernelBuilderPtr AicpuKernelBuilder::Instance()
{
    static std::once_flag flag;
    std::call_once(flag, [&]() {
        instance_.reset(new (std::nothrow) AicpuKernelBuilder);
    });
    return instance_;
}

ge::Status AicpuKernelBuilder::CalcOpRunningParam(const ge::Node &node) const
{
    std::string node_name = node.GetName();
    std::string node_type = node.GetType();
    AICPUE_LOGI("AICPUKernel's op %s[%s] run CalcOpRunningParam", node_name.c_str(), node_type.c_str());

    std::shared_ptr<ge::OpDesc> op_desc_ptr = node.GetOpDesc();
    AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, ErrorCode::INPUT_PARAM_NULL)

    // workspace size is 0
    int64_t total_size = 0;
    op_desc_ptr->SetWorkspaceBytes({total_size});
    AICPU_CHECK_RES_WITH_LOG(KernelBuilder::SetOutPutsSize(op_desc_ptr),
        "Call SetOutPutsSize function failed, op[%s].", node_name.c_str());

    // Set workspace memory reuse flag
    AICPU_CHECK_FALSE_EXEC(ge::AttrUtils::SetListBool(op_desc_ptr, kWorkspaceReuseFlag, {false}),
        AICPU_REPORT_INNER_ERROR("Set workspace memory reuse flag failed, op[%s].",
            node_name.c_str());
        return ErrorCode::ADD_ATTR_FAILED)

    AICPUE_LOGI("AICPUKernel's op %s[%s] run CalcOpRunningParam success. workspace total size is %lld",
                node_name.c_str(), node_type.c_str(), total_size);
    return ge::SUCCESS;
}

ge::Status AicpuKernelBuilder::GenerateTask(const ge::Node &node, const ge::RunContext &run_context, std::vector<domi::TaskDef> &tasks)
{
    AICPUE_LOGI("AICPUKernel's op %s[%s] run GenerateTask.", node.GetName().c_str(), node.GetType().c_str());
    ge::OpDescPtr op_desc_ptr = node.GetOpDesc();
    AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, INPUT_PARAM_NULL)
    AICPU_CHECK_NOTNULL_ERRCODE(run_context.model, INPUT_PARAM_NULL)
    AICPU_CHECK_NOTNULL_ERRCODE(run_context.stream, INPUT_PARAM_NULL)

    domi::TaskDef task_def;
    task_def.set_type(RT_MODEL_TASK_KERNEL);
    // no need to set streamID for task_def, ge will reallocatee stream
    domi::KernelDef *kernel_def = task_def.mutable_kernel();
    AICPU_CHECK_NOTNULL(kernel_def)
    // create prame_base and launch kernel
    AICPU_CHECK_RES_WITH_LOG(BuildAndLaunchKernel(op_desc_ptr, run_context, kernel_def),
        "Call BuildAndLaunchKernel function failed, op[%s].", node.GetName().c_str())

    tasks.emplace_back(task_def);
    AICPUE_LOGI("AICPUKernel's op %s[%s] run GenerateTask success.", node.GetName().c_str(), node.GetType().c_str());
    return ge::SUCCESS;
}

// Make and init task extend info
ge::Status AicpuKernelBuilder::MakeAicpuKernelExtInfo(const ge::OpDescPtr &op_desc_ptr,
                                                  vector<char> &task_ext_info) const
{
    AICPU_CHECK_RES_WITH_LOG(MakeBaseExtInfo(op_desc_ptr, task_ext_info),
        "Call MakeTaskExtInfo funtion failed, op[%s].", op_desc_ptr->GetName().c_str())
    // calc ExtInfo length
    uint64_t extend_info_len = task_ext_info.size();

    extend_info_len += aicpu::FWKAdapter::kExtInfoHeadSize;
    extend_info_len += sizeof(SessionInfo);

    uint64_t ext_info_offset = task_ext_info.size();
    task_ext_info.resize(extend_info_len, 0);
    char *ext_info_buf = task_ext_info.data();

    // init ext info 1: input ShapeAndType
    auto ext_info = reinterpret_cast<aicpu::FWKAdapter::ExtInfo *>(ext_info_buf + ext_info_offset);
    ext_info->infoType = aicpu::FWKAdapter::FWK_ADPT_EXT_SESSION_INFO;
    ext_info->infoLen = (sizeof(SessionInfo));
    ext_info_offset += aicpu::FWKAdapter::kExtInfoHeadSize;
    SessionInfo *session_info = reinterpret_cast<SessionInfo *>(ext_info_buf + ext_info_offset);
    session_info->sessionId = GenerateUniqueSessionId();
    session_info->kernelId = GenerateUniqueKernelId();
    session_info->sessFlag = false;
    return ge::SUCCESS;
}

ge::Status AicpuKernelBuilder::BuildAndLaunchKernel(
    const ge::OpDescPtr &op_desc_ptr, const ge::RunContext &run_context,
    domi::KernelDef *&kernel_def) const
{
    std::vector<char> task_ext_info;
    AICPU_CHECK_RES_WITH_LOG(MakeAicpuKernelExtInfo(op_desc_ptr, task_ext_info),
        "Call MakeAicpuKernelExtInfo function failed, op[%s]", op_desc_ptr->GetName().c_str())
    if (task_ext_info.size() == 0) {
        kernel_def->set_kernel_ext_info_size(0);
    } else {
        kernel_def->set_kernel_ext_info(reinterpret_cast<void *>(task_ext_info.data()), task_ext_info.size());
        kernel_def->set_kernel_ext_info_size(task_ext_info.size());
    }

    // set kernel so name's len and offset addr
    std::string kernel_so_name;
    (void)ge::AttrUtils::GetStr(op_desc_ptr, "kernelSo", kernel_so_name);
    if (kernel_so_name.empty()){
        kernel_so_name = kDefaultKernelSo;
    }
    // set kernel function name's len and offset addr
    std::string func_name;
    (void)ge::AttrUtils::GetStr(op_desc_ptr, "funcName", func_name);
    if (func_name.empty()){
        func_name = kDefaultFunctionName;
    }
    // calculate param_len: AicpuParamHead.len + io_addrs_size + notify_id.len + customizedAttr.len
    bool flag_async = false;
    (void)ge::AttrUtils::GetBool(op_desc_ptr, "paramLen", flag_async);
    uint32_t param_len = flag_async ? static_cast<uint32_t>(sizeof(AicpuParamHead) + sizeof(uint32_t)) :
                                           static_cast<uint32_t>(sizeof(AicpuParamHead));
    // get input and output total number, no need to check overflow
    uint32_t io_addrs_num = static_cast<uint32_t>(op_desc_ptr->GetInputsSize() + op_desc_ptr->GetOutputsSize());
    // get input and output addrs size, no need to check overflow
    uint32_t io_addrs_size = io_addrs_num * static_cast<uint32_t>(sizeof(uint64_t));
    // refresh param_len, no need to check overflow
    param_len += io_addrs_size;

    ge::GeAttrValue::BYTES bytes;
    bool has_customized_attr = ge::AttrUtils::GetZeroCopyBytes(op_desc_ptr, kCustomizedOpDef, bytes);
    // When it's aicpu customized ops, get customized attr
    if (has_customized_attr) {
        CHECK_UINT32_ADD_OVERFLOW(param_len, sizeof(uint32_t), ErrorCode::DATA_OVERFLOW,
            "Overflow when param total bytes[%u] add 4bytes, op[%s]",
            param_len, op_desc_ptr->GetName().c_str())
        param_len += sizeof(uint32_t);
        // Customized attr length must be less than UINT32_MAX, no need to check overflow
        uint32_t customized_attr_len = static_cast<uint32_t>(bytes.GetSize());
        CHECK_UINT32_ADD_OVERFLOW(param_len, customized_attr_len, ErrorCode::DATA_OVERFLOW,
            "Overflow when calculate total bytes of task param[%u] and custom "
            "attr[%u], op[%s]", param_len, customized_attr_len,
            op_desc_ptr->GetName().c_str())
        param_len += customized_attr_len;
    }

    uint32_t ext_info_length = task_ext_info.size();
    uint64_t ext_info_addrs = 0;
    std::string task_args;
    // Create task_args: AicpuParamHead + ioAddrs + notify_id + customizedAttr
    AicpuParamHead param_head = {param_len, io_addrs_num, ext_info_length, ext_info_addrs};
    task_args.append(reinterpret_cast<const char *>(&param_head), sizeof(AicpuParamHead));
    // TaskArgs append ioAddrs
    task_args.append(io_addrs_size, ' ');
    // TaskArgs append notify_id for Async op
    if (flag_async) {
        uint32_t notify_id = 0;
        task_args.append(reinterpret_cast<const char *>(&notify_id), sizeof(uint32_t));
    }

    // When it's aicpu customized ops, task_args should append customized attr
    if (has_customized_attr) {
        const uint8_t *node_def_data = bytes.GetData();
        AICPU_CHECK_NOTNULL(node_def_data)
        size_t customized_attr_len = bytes.GetSize();
        uint32_t node_def_len = static_cast<uint32_t>(customized_attr_len);
        task_args.append(reinterpret_cast<const char *>(&node_def_len), sizeof(uint32_t));
        task_args.append(reinterpret_cast<const char *>(node_def_data), customized_attr_len);
    }

    kernel_def->set_args(task_args.data(), param_len);
    kernel_def->set_args_size(param_len);
    kernel_def->set_so_name(kernel_so_name);
    kernel_def->set_kernel_name(func_name);

    // get kernel type
    bool cust_aicpu_flag = false;
    (void) ge::AttrUtils::GetBool(op_desc_ptr, kCustAicpuFlag, cust_aicpu_flag);
    cce::ccKernelType kernel_type = cust_aicpu_flag ? cce::ccKernelType::CUST_AI_CPU : cce::ccKernelType::AI_CPU;

    domi::KernelContext *context = kernel_def->mutable_context();
    AICPU_CHECK_NOTNULL(context)
    context->set_op_index(op_desc_ptr->GetId());
    context->set_kernel_type(kernel_type);

    AICPUE_LOGI("Task_args length is %zu, param_len is %u, kernel_so_name is %s, func_name is %s, kernel_type is %d",
                task_args.length(), param_len, kernel_so_name.c_str(), func_name.c_str(), kernel_type);
    return ge::SUCCESS;
}

FACTORY_KERNEL_BUILDER_CLASS_KEY(AicpuKernelBuilder, "AICPUBuilder")
}
